import datetime
import json
import logging
import os

import pytz
from absl import app, flags
import google.protobuf.json_format as json_format

from coin.proto.coin_strategy_summary_pb2 import StrategyStatLog 
from coin.support.pta.logging.strategy_stat_logger import StrategyStatReader


FLAGS = flags.FLAGS


class OnLog(object):
  def __init__(self, symbol_filter):
    self._symbol_filter = symbol_filter
    self.record_list = []

  def _log_filtered(self, data):
    if self._symbol_filter is not None:
      if data.get('symbol') not in self._symbol_filter:
        return True
    return False

  def on_log(self, timestamp, log):
    pb = StrategyStatLog()
    pb.ParseFromString(log)
    data = json_format.MessageToDict(pb, preserving_proto_field_name=True)
    if self._log_filtered(data):
      return
    hr_ts = datetime.datetime.fromtimestamp(timestamp / 1e9).strftime('%Y%m%d %H:%M:%S.%f')
    data['record_timestamp'] = timestamp
    data['record_timestamp_human_readable'] = hr_ts
    self.record_list.append(data)


def main(_):
  root_dir = os.path.expanduser(FLAGS.root_dir)
  machine = FLAGS.machine
  strategy_name = FLAGS.strategy_name
  start_time = FLAGS.start_time
  end_time = FLAGS.end_time
  symbol_filter = FLAGS.symbol
  assert machine, '--machine must be specified.'
  assert strategy_name, '--strategy_name must be specified.'
  assert start_time is not None and end_time is not None, (start_time, end_time)

  dt_format = '%Y%m%dT%H%M%S'
  start_dt = datetime.datetime.strptime(
      start_time, dt_format).replace(tzinfo=pytz.UTC)
  end_dt = datetime.datetime.strptime(
      end_time, dt_format).replace(tzinfo=pytz.UTC)

  if symbol_filter is not None:
    symbol_filter = [elem.strip() for elem in symbol_filter.split(',')]

  start_dt_str = start_dt.strftime(dt_format)
  end_dt_str = end_dt.strftime(dt_format)
  print('Running for %s-%s %s ...' % (start_dt_str, end_dt_str, machine))
  onlog = OnLog(symbol_filter=symbol_filter)
  stat_reader = StrategyStatReader()
  stat_reader.run_from_archive(
      on_log_callback=onlog.on_log,
      start_time=start_dt,
      end_time=end_dt,
      root_dir=root_dir,
      machine=machine,
      strategy_name=strategy_name)
  filename = 'strat_stat_log.%s-%s.%s.json' % (start_dt_str, end_dt_str, strategy_name)
  with open(filename, 'w') as out_file:
    if FLAGS.unpretty:
      records = '[\n' + \
                ',\n'.join(json.dumps(record) for record in onlog.record_list) + \
                '\n]'
      out_file.write(records)
    else:
      json.dump(onlog.record_list, out_file, indent=2)
  print(filename)
  


if __name__ == '__main__':
  flags.DEFINE_string('root_dir', '~/data/strat_proto_log', 'root_dir.')
  flags.DEFINE_string('machine', None, 'Instance machine name.')
  flags.DEFINE_string('strategy_name', None, 'Strategy name.')
  flags.DEFINE_string('start_time', None, '%Y%m%dT%H%M%S')
  flags.DEFINE_string('end_time', None, '%Y%m%dT%H%M%S')
  flags.DEFINE_string('symbol', None, '')
  flags.DEFINE_bool('unpretty', False, '')
  logging.basicConfig(
      level='DEBUG',
      format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
